# encoding:utf-8
from urllib.parse import urlparse

import psycopg2
from llama_index.core import StorageContext, VectorStoreIndex, Document, ServiceContext, SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from llama_index.vector_stores.postgres import PGVectorStore
# from setuptools.config.expand import read_files
from sqlalchemy import make_url

from config import PGVECTOR_URL, VECTOR_DB_NAME, VECTOR_TABLE_NAME
from src.agents.llama_index_models import embed_model
db_name = VECTOR_DB_NAME
table_name = VECTOR_TABLE_NAME
connection_string =PGVECTOR_URL
url = make_url(connection_string)
vector_store = PGVectorStore.from_params(
    database=db_name,
    host=url.host,
    password=url.password,
    port=url.port,
    user=url.username,
    table_name=table_name,
    embed_dim=1792,  # openai embedding dimension
    hybrid_search=True,
    text_search_config="chinese",
    hnsw_kwargs={
        "hnsw_m": 16,
        "hnsw_ef_construction": 64,
        "hnsw_ef_search": 40,
        "hnsw_dist_method": "vector_cosine_ops",
    },
)

def storage_documents():
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    # 这里要改
    with open(r"D:\project\doc-ai\src\agents\gmp规范.txt","r",encoding="utf-8") as f:
        file_contents = f.read()

    documents = Document(text=file_contents,metadata={"filename":"药品生产质量管理规范(GMP).txt","file_id":"444"})
    index = VectorStoreIndex.from_documents(
        [documents], storage_context=storage_context, show_progress=True,
        embed_model=embed_model,
        transformations=[SentenceSplitter(chunk_size=2048, chunk_overlap=100)],
    )
#
# def storage_docx():
#     storage_context = StorageContext.from_defaults(vector_store=vector_store)
#     reader = SimpleDirectoryReader(
#         input_dir=r"D:\project\doc-ai\src\agents\docs\发泓睿-公用操作规程--2024.08.28\生产过程管理"
#     )
#     documents = reader.load_data()
#     for document in documents:
#         document.text = document.text.replace('\x00', '')
#     index = VectorStoreIndex.from_documents(
#         documents, storage_context=storage_context, show_progress=True,
#         embed_model=embed_model,
#         transformations=[SentenceSplitter(chunk_size=2048, chunk_overlap=100)],
#     )


def storage_api(file_contents, metadata):
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    '''存储前端 信息'''
    documents = Document(text=file_contents, metadata=metadata)
    print(f"存储的信息:{documents}")
    index = VectorStoreIndex.from_documents(
        [documents], storage_context=storage_context, show_progress=True,
        embed_model=embed_model,
        transformations=[SentenceSplitter(chunk_size=1024, chunk_overlap=100)],
    )
    print(f"存储地址：{index}")

def delete_files(belong_file_id):
    # 解析数据库连接 URL
    result = urlparse(PGVECTOR_URL)
    username = result.username
    password = result.password
    hostname = result.hostname
    port = result.port

    # 连接数据库
    connection = psycopg2.connect(
        database=db_name,
        user=username,
        password=password,
        host=hostname,
        port=port
    )
    cur = connection.cursor()

    try:
        # 动态生成 SQL 语句
        sql = f"""
            DO $$
            BEGIN
                -- 检查表是否存在
                IF EXISTS (
                    SELECT 1
                    FROM information_schema.tables
                    WHERE table_schema = 'public'
                    AND table_name = 'data_{table_name}'
                ) THEN
                    -- 如果表存在，检查是否有符合条件的记录
                    IF EXISTS (
                        SELECT 1
                        FROM public.data_{table_name}
                        WHERE metadata_->>'belong_file_id' = '{belong_file_id}'
                    ) THEN
                        -- 如果有符合条件的记录，执行删除操作
                        DELETE FROM public.data_{table_name}
                        WHERE metadata_->>'belong_file_id' = '{belong_file_id}';
                    END IF;
                END IF;
            END $$;
            """

        # 执行 SQL 语句
        cur.execute(sql)
        connection.commit()
        print(f"Deleted records with belong_file_id={belong_file_id} from table data_{table_name}")

    except Exception as e:
        # 如果发生错误，回滚事务
        connection.rollback()
        print(f"Error occurred: {e}")
    finally:
        # 关闭游标和连接
        cur.close()
        connection.close()


def delete_konwledge(parent_knowledge_id):
    # 解析数据库连接 URL
    result = urlparse(PGVECTOR_URL)
    username = result.username
    password = result.password
    hostname = result.hostname
    port = result.port

    # 连接数据库
    connection = psycopg2.connect(
        database=db_name,
        user=username,
        password=password,
        host=hostname,
        port=port
    )
    cur = connection.cursor()

    try:
        # 动态生成 SQL 语句
        sql = f"""
                DO $$
                BEGIN
                    -- 检查表是否存在
                    IF EXISTS (
                        SELECT 1
                        FROM information_schema.tables
                        WHERE table_schema = 'public'
                        AND table_name = 'data_{table_name}'
                    ) THEN
                        -- 如果表存在，检查是否有符合条件的记录
                        IF EXISTS (
                            SELECT 1
                            FROM public.data_{table_name}
                            WHERE metadata_->>'parent_knowledge_id' = '{parent_knowledge_id}'
                        ) THEN
                            -- 如果有符合条件的记录，执行删除操作
                            DELETE FROM public.data_{table_name}
                            WHERE metadata_->>'parent_knowledge_id' = '{parent_knowledge_id}';
                        END IF;
                    END IF;
                END $$;
                """

        # 执行 SQL 语句
        cur.execute(sql)
        connection.commit()

    except Exception as e:
        # 如果发生错误，回滚事务
        connection.rollback()
        print(f"Error occurred: {e}")
    finally:
        # 关闭游标和连接
        cur.close()
        connection.close()


if __name__ == "__main__":
    # storage_docx()
    # pass
    delete_files("123")
